#!/bin/bash
pwd=$(cd `dirname $0`;pwd)
echo $pwd
cd $pwd/cann-ops-adv
#cd ops_adv
date=$(date +%Y%m%d)
startTime=`date +%Y%m%d-%H:%M:%S`
startTime_s=`date +%s`
source /usr/local/Ascend/ascend-toolkit/set_env.sh

rm -rf $pwd/result_fa
mkdir -p $pwd/result_fa
result=Failed
result_FA=Failed
result_unpadFA=Failed
result_fag=Failed
result_unpadfag=Failed

FFN_overall_result=Failed
result_ffn_v2_antiquant=Failed
result_ffn_v3_quant_token_index_flag=Failed
result_ffn_v3_float16_high_precision=Failed

PFA_overall_result=Failed
result_PFA=Failed
result_PFA_v2=Failed
result_PFA_v3=Failed

FIA_overall_result=Failed
result_FIA=Failed
result_FIA_v2=Failed

IFA_overall_result=Failed
result_IFA=Failed
result_IFA_v2=Failed
result_IFA_v3=Failed
result_IFA_v4=Failed


#bash build.sh -e > $pwd/result_fa/result_fa.log 2>&1
bash build.sh -e --disable-check-compatible | tee $pwd/result_fa/result_fa.log 2>&1

count_fag=$(grep -rn "Run test_flash_attention_score_grad success" $pwd/result_fa/result_fa.log |wc -l)
count_unpadfa=$(grep -rn "Run test_flash_attention_varLen_score success" $pwd/result_fa/result_fa.log |wc -l)
count_FA=$(grep -rn "Run test_flash_attention_score success" $pwd/result_fa/result_fa.log |wc -l)
count_unpadfag=$(grep -rn "Run test_flash_attention_unpadding_score_grad success" $pwd/result_fa/result_fa.log |wc -l)

count_ffn_v2_antiquant=$(grep -rn "Run ffn_v2_antiquant success" $pwd/result_fa/result_fa.log |wc -l)
count_ffn_v3_quant_token_index_flag=$(grep -rn "Run ffn_v3_quant_token_index_flag success" $pwd/result_fa/result_fa.log |wc -l)
count_ffn_v3_float16_high_precision=$(grep -rn "Run ffn_v3_float16_high_precision success" $pwd/result_fa/result_fa.log |wc -l)

count_PFA=$(grep -rn "Run test_prompt_flash_attention success" $pwd/result_fa/result_fa.log |wc -l)
count_PFA_v2=$(grep -rn "Run test_prompt_flash_attention_v2 success" $pwd/result_fa/result_fa.log |wc -l)
count_PFA_v3=$(grep -rn "Run test_prompt_flash_attention_v3 success" $pwd/result_fa/result_fa.log |wc -l)

count_FIA=$(grep -rn "Run test_fused_infer_attention_score success" $pwd/result_fa/result_fa.log |wc -l)
count_FIA_v2=$(grep -rn "Run test_fused_infer_attention_score_v2 success" $pwd/result_fa/result_fa.log |wc -l)

count_IFA=$(grep -rn "Run test_incre_flash_attention success" $pwd/result_fa/result_fa.log |wc -l)
count_IFA_v2=$(grep -rn "Run test_incre_flash_attention_v2 success" $pwd/result_fa/result_fa.log |wc -l)
count_IFA_v3=$(grep -rn "Run test_incre_flash_attention_v3 success" $pwd/result_fa/result_fa.log |wc -l)
count_IFA_v4=$(grep -rn "Run test_incre_flash_attention_v4 success" $pwd/result_fa/result_fa.log |wc -l)


if [ $count_FA -eq 1 ];then
        result_FA=Pass
fi

if [ $count_unpadfa -eq 1 ];then
        result_unpadFA=Pass
fi

if [ $count_fag -eq 1 ];then
        result_fag=Pass
fi

if [ $count_unpadfag -eq 1 ];then
        result_unpadfag=Pass
fi

if [ "${result_FA}" == "Pass" -a "${result_unpadFA}" == "Pass" -a "${result_fag}" == "Pass" -a "${result_unpadfag}" == "Pass" ];then
            result=Pass
fi

if [ $count_ffn_v2_antiquant -eq 1 ];then
        result_ffn_v2_antiquant=Pass
fi

if [ $count_ffn_v3_quant_token_index_flag -eq 1 ];then
        result_ffn_v3_quant_token_index_flag=Pass
fi

if [ $count_ffn_v3_float16_high_precision -eq 1 ];then
        result_ffn_v3_float16_high_precision=Pass
fi

if [ "${result_ffn_v2_antiquant}" == "Pass" -a "${result_ffn_v3_quant_token_index_flag}" == "Pass" -a "${result_ffn_v3_float16_high_precision}" == "Pass" ];then
            FFN_overall_result=Pass
fi

if [ $count_PFA -eq 1 ];then
    result_PFA=Pass
fi

if [ $count_PFA_v2 -eq 1 ];then
    result_PFA_v2=Pass
fi

if [ $count_PFA_v3 -eq 1 ];then
    result_PFA_v3=Pass
fi

if [ "${result_PFA}" == "Pass" -a "${result_PFA_v2}" == "Pass" -a "${result_PFA_v3}" == "Pass" ];then
            PFA_overall_result=Pass
fi

if [ $count_FIA -eq 1 ];then
    result_FIA=Pass
fi

if [ $count_FIA_v2 -eq 1 ];then
    result_FIA_v2=Pass
fi

if [ "${result_FIA}" == "Pass" -a "${result_FIA_v2}" == "Pass" ];then
            FIA_overall_result=Pass
fi

if [ $count_IFA -eq 1 ];then
    result_IFA=Pass
fi

if [ $count_IFA_v2 -eq 1 ];then
    result_IFA_v2=Pass
fi

if [ $count_IFA_v3 -eq 1 ];then
    result_IFA_v3=Pass
fi

if [ $count_IFA_v4 -eq 1 ];then
    result_IFA_v4=Pass
fi

if [ "${result_IFA}" == "Pass" -a "${result_IFA_v2}" == "Pass" -a "${result_IFA_v3}" == "Pass" -a "${result_IFA_v4}" == "Pass" ];then
            IFA_overall_result=Pass
fi

echo "test_flash_attention_score: ${result_FA}"  | tee $pwd/result_fa/case_result.txt
echo "test_flash_attention_varLen_score: ${result_unpadFA}" | tee -a $pwd/result_fa/case_result.txt
echo "test_flash_attention_score_grad: ${result_fag}" | tee -a $pwd/result_fa/case_result.txt
echo "test_flash_attention_unpadding_score_grad: ${result_unpadfag}" | tee -a $pwd/result_fa/case_result.txt
echo "test_ffn_v2_antiquant: ${result_ffn_v2_antiquant}" | tee -a $pwd/result_fa/case_result.txt
echo "test_ffn_v3_quant_token_index_flag: ${result_ffn_v3_quant_token_index_flag}" | tee -a $pwd/result_fa/case_result.txt
echo "test_ffn_v3_float16_high_precision: ${result_ffn_v3_float16_high_precision}" | tee -a $pwd/result_fa/case_result.txt
echo "test_PFA: ${result_PFA}" | tee -a $pwd/result_fa/case_result.txt
echo "test_PFA_v2: ${result_PFA_v2}" | tee -a $pwd/result_fa/case_result.txt
echo "test_PFA_v3: ${result_PFA_v3}" | tee -a $pwd/result_fa/case_result.txt
echo "test_FIA: ${result_FIA}" | tee -a $pwd/result_fa/case_result.txt
echo "test_FIA_v2: ${result_FIA_v2}" | tee -a $pwd/result_fa/case_result.txt
echo "test_IFA: ${result_IFA}" | tee -a $pwd/result_fa/case_result.txt
echo "test_IFA_v2: ${result_IFA_v2}" | tee -a $pwd/result_fa/case_result.txt
echo "test_IFA_v3: ${result_IFA_v3}" | tee -a $pwd/result_fa/case_result.txt
echo "test_IFA_v4: ${result_IFA_v4}" | tee -a $pwd/result_fa/case_result.txt
echo "FA Result: ${result}" | tee -a $pwd/result_fa/case_result.txt
echo "FFN Result: ${FFN_overall_result}" | tee -a $pwd/result_fa/case_result.txt
echo "PFA Result: ${PFA_overall_result}" | tee -a $pwd/result_fa/case_result.txt
echo "FIA Result: ${FIA_overall_result}" | tee -a $pwd/result_fa/case_result.txt
echo "IFA Result: ${IFA_overall_result}" | tee -a $pwd/result_fa/case_result.txt
echo "---------------------------------------------------------"

if grep -q "Failed" $pwd/result_fa/case_result.txt;then
    echo "FASmokeTest: Failed" | tee $pwd/result_fa/case_result.txt
else
    echo "FASmokeTest: Pass" | tee $pwd/result_fa/case_result.txt
fi
cp $pwd/result_fa/case_result.txt $pwd/result.txt
endTime=`date +%Y%m%d-%H:%M:%S`
endTime_s=`date +%s`
sumTime=$[ $endTime_s - $startTime_s ]
echo -e "-------------------\033[32m 端到端执行用例耗时:$sumTime seconds!!!----------"
echo "---------------------------------------------------------"
